#include "EQWave.h"
#include <iostream>
#include <mgl2/mgl.h>

using namespace arma;
using namespace std;

void EQWave::set_a(double diff_volecity)
{
	a = diff_volecity;
}

void EQWave::set_xlimits(double lower, double upper)
{
	xlimits[0] = lower;
	xlimits[1] = upper;
}
	
void EQWave::set_tlimits(double lower, double upper)
{
	tlimits[0] = lower;
	tlimits[1] = upper;
}
	
void EQWave::set_BC_left(f1v BC_left)
{
	_BC_left = BC_left;
}

void EQWave::set_BC_left(double BC_left)
{
	BC_L = BC_left;
}

void EQWave::set_BC_right(f1v BC_right)
{
	_BC_right = BC_right;
}

void EQWave::set_BC_right(double BC_right)
{
	BC_R = BC_right;
}

void EQWave::set_IC1(f1v IC1)
{
	_IC1 = IC1;
}

void EQWave::set_IC2(f1v IC2)
{
	_IC2 = IC2;
}

void EQWave::set_hx(double h)
{
	hx = h;
}

void EQWave::set_ht(double tau)
{
	ht = tau;
}


unsigned int EQWave::compute_Nx()
{
	Nx = (xlimits[1] - xlimits[0])/hx;
	return Nx;
}
	
unsigned int EQWave::compute_Nt()
{
	Nt = (tlimits[1] - tlimits[0])/ht;
	return Nt;
}

void EQWave::init()
{
	if(a*ht >= hx)
	{
		cout << "Error, must meet the condition: a*ht < hx !!! \n";
		exit(1);
	}

	d = a*a*ht*ht/hx/hx;
	compute_Nx();
	compute_Nt();
	u.resize(Nt+1);

	for (int i = 0; i < Nt+1; ++i)
	{
		u[i].set_size(Nx+1);
		u[i].zeros();
	}
	
	for (int j = 0; j < Nx+1; ++j)
	{
		u[0](j) = (*_IC1)(xlimits[0]+j*hx);
	}

	
	/* for the DOULE BC CASE */
	for (int i = 0; i < Nt+1; ++i)
	{
		u[i](0) = (*_BC_left)(tlimits[0] + i*ht);
		u[i](Nx) = (*_BC_right)(tlimits[0] + i*ht);
	}

	for (int j = 1; j < Nx; ++j)
	{
		u[1](j) = (*_IC2)(xlimits[0] + j*hx)*ht + u[0](j) + d*(u[0](j+1)-2*u[0](j) + u[0](j-1))/2;
	}	
}

void EQWave::solve()
{
	init();
	for (int i = 2; i < Nt+1; ++i)
	{
		for (int j = 1; j < Nx; ++j)
		{
			u[i](j) = 2*u[i-1](j) - u[i-2](j) + d*(u[i-1](j+1)+u[i-1](j-1)-2*u[i-1](j));
		}
	}
}

void EQWave::draw_data(const char* filename, std::vector<double> time) const
{
	const char* lineScheme[10] = {"p-1", "r|1", "g;1", "b:1", "cj1", "mi1", "y=1", "l1", "e1", "k1"};

	mglData mdata(Nx+1), t(Nx+1);
	for (int i = 0; i < Nx+1; ++i)
	{
		t.a[i] = tlimits[0] + i*hx;
	}
	double ymin = 0, ymax = 0;
	// for (int i = 0; i < time.size(); ++i)
	// {
	// 	if()
	// }
	mglGraph gr;
	gr.Alpha(true);
	gr.Light(true);
	for (int i = 0; i < time.size(); ++i)
	{
		int ind = (time[i]-tlimits[0])/ht;
		if (u[ind].max() > ymax)
			ymax = u[ind].max();
		if (u[ind].min() < ymin)
			ymin = u[ind].min();
	}
	gr.SetRanges(xlimits[0],xlimits[1], floor(ymin), ceil(ymax));
	// gr.Box();
	gr.Axis();
	gr.Grid("xy", "W");
	gr.Label('x',"x",1);
	gr.Label('y',"u(x,t)",1);
	char legend[256];
	double *a = new double[Nx+1];
	for (int i = 0; i < time.size(); ++i)
	{
		int ind = (time[i]-tlimits[0])/ht;
		for (int j = 0; j < Nx+1; ++j)
		{			
			a[j] = u[ind](j);
		}
		mdata.Set(a, Nx+1);
		sprintf(legend, "t = %f", time[i]);
		gr.Plot(t, mdata, lineScheme[i]);
		gr.AddLegend(legend,  lineScheme[i]);
	}

	gr.Legend(3,"A#");
	gr.WriteFrame(filename);
	delete []a;
}

void EQWave::draw_data(const char* filename, double time_length) const
{
	mglData mdata(Nx+1), t(Nx+1);
	for (int i = 0; i < Nx+1; ++i)
	{
		t.a[i] = tlimits[0] + i*hx;
	}
	double ymin = 0, ymax = 0;
	int length = (time_length-tlimits[0])/ht;
	int delta_L = 1;
	if (length > 200)
		delta_L = length / 200;
	for (int i = 0; i < length; i += delta_L)
	{
		if (u[i].max() > ymax)
			ymax = u[i].max();
		if (u[i].min() < ymin)
			ymin = u[i].min();
	}
	// for (int i = 0; i < time.size(); ++i)
	// {
	// 	if()
	// }
	// cout << length << endl;
	mglGraph gr;
	gr.StartGIF(filename);
	double *a = new double[Nx+1];
	for (int i = 0; i < length; i += delta_L)
	{
		for (int j = 0; j < Nx+1; ++j)
		{			
			a[j] = u[i](j);
		}
		gr.NewFrame();
		gr.Alpha(true);
		gr.Light(true);
		gr.SetRanges(xlimits[0],xlimits[1], floor(ymin), ceil(ymax));
		// gr.Box();
		gr.Axis();
		gr.Grid("xy", "W");
		gr.Label('x',"x",1);
		gr.Label('y',"u(x,t)",1);
		mdata.Set(a, Nx+1);
		gr.Plot(t, mdata);
		gr.EndFrame();
	}
	gr.CloseGIF();
	delete []a;
}

void EQWave::save_data(const char* filename) const
{
	using namespace std;

	ofstream fout;
	fout.open(filename);
	//fout.precision(4);
	//fout.width(5);
	for (int i = 0; i < Nt+1; ++i)
	{
		fout << "# t = " << tlimits[0] + i*ht << endl;
		for (int j = 0; j < Nx; ++j)
		{
			fout << u[i](j) << ", ";			
		}
		fout << u[i](Nx) << ";" << endl << endl;
	}
}